#include <hal/ahci/ahci.h>
#include <hal/ahci/hba.h>
#include <hal/ahci/sata.h>
#include <hal/pci.h>

#include <arch/io.h>
#include <arch/vectors.h>
#include <arch/interrupts.h>

#include <mm/mm.h>
#include <mm/mmio.h>
#include <mm/valloc.h>

#include <device/block.h>
#include <device/blkio.h>

#include <common.h>
#include <string.h>

#define LOG_TAG "AHCI"
#include <log.h>

hba_t hba;

void ahci_blkio_handler(blkio_req_t* req);

extern void ahci_hba_isr(const isr_param* param);

static int get_free_slot(hba_port_t* port) {
    u32 sact = port->regs[HBA_REG_PxSACT];
    u32 ci = port->regs[HBA_REG_PxCI];
    u32 free_bmap = sact | ci;
    int idx = 0;
    while (idx <= (int)hba.cmd_slots && (free_bmap & 0x1)) {
        idx++;
        free_bmap >>= 1;
    }

    if (idx > (int)hba.cmd_slots) idx = -1;
    return idx;
}

int hba_prepare_cmd(hba_port_t* port, hba_cmdt_t** cmdt, hba_cmdh_t** cmdh) {
    int slot = get_free_slot(port);
    if (slot == -1) {
        LOGE("No free command slot\n");
    }

    *cmdh = &port->cmdlst[slot];
    *cmdt = vzalloc_dma(sizeof(hba_cmdt_t));

    memset(*cmdh, 0, sizeof(hba_cmdh_t));

    (*cmdh)->cmd_table_base = get_pa((uptr)(*cmdt));
    (*cmdh)->options = HBA_CMDH_FIS_LEN(sizeof(sata_cmd_fis_t)) | HBA_CMDH_CLR_BUSY;

    return slot;
}

void ahci_register_device(struct hba_device* hbadev) {
    blk_dev_t* bdev = block_alloc_dev(hbadev->model, hbadev, ahci_blkio_handler);

    bdev->end_lba = hbadev->max_lba;
    bdev->blk_size = hbadev->block_size;

    block_mount(bdev);
}

void ahci_init() {
    pci_dev_t* dev = pci_get_by_class(AHCI_HBA_CLASS);
    if (!dev) {
        LOGE("No AHCI device found\n");
        return;
    }

    u32 bar, size;
    size = pci_bar_sizing(dev, &bar, 6);

    if (!PCI_BAR_MMIO(bar)) {
        LOGE("BAR#6 is not MMIO\n");
        return;
    }
    u32 cmd = pci_read_cspace(dev->cspace_base, PCI_REG_STATUS_CMD);
    cmd |= PCI_CMD_DISABLE_INTR | PCI_CMD_MM_ACCESS | PCI_CMD_BUS_MASTER;

    pci_write_cspace(dev->cspace_base, PCI_REG_STATUS_CMD, cmd);

    pci_setup_msi(dev, AHCI_HBA_IV);
    intr_register_isr(AHCI_HBA_IV, ahci_hba_isr);

    memset(&hba, 0, sizeof(hba));
    hba.base = (uptr*)ioremap(PCI_BAR_ADDR_MM(bar), size >> 12);

    hba.base[HBA_REG_GHC] |= HBA_GHC_ACHI_ENABLE;
    hba.base[HBA_REG_GHC] |= HBA_GHC_INTR_ENABLE;

    u32 cap = hba.base[HBA_REG_CAP];
    u32 pmap = hba.base[HBA_REG_PI];

    hba.ports_num = (cap & 0x1F) + 1;
    hba.cmd_slots = (cap >> 8) & 0x1F;
    hba.version = hba.base[HBA_REG_VS];
    hba.ports_bmap = pmap;

    /* Configure HBA ports */
    uptr clb_pg_va, fis_pg_va, clb_pa, fis_pa;
    for (size_t i = 0, fis_i = 0, clb_i = 0;
        i < 32;
        i++, pmap >>= 1, fis_i = (fis_i + 1) % 16, clb_i = (clb_i + 1) % 4
        ) {
        if (!(pmap & 0x1)) {
            continue;
        }

        hba_port_t* port = valloc(sizeof(hba_port_t));
        u32* port_regs = (u32*)(&hba.base[HBA_REG_PBASE + i * HBA_REG_PSIZE]);

        hba_reset_port(port_regs);

        if (!clb_i) {
            // 每页最多4个命令队列
            clb_pa = get_page(PG_LOCKED);
            clb_pg_va = ioremap(clb_pa, 1);
            memset((void*)clb_pg_va, 0, PG_SIZE);
        }
        if (!fis_i) {
            // 每页最多16个FIS
            fis_pa = get_page(PG_LOCKED);
            fis_pg_va = ioremap(fis_pa, 1);
            memset((void*)fis_pg_va, 0, PG_SIZE);
        }

        port_regs[HBA_REG_PxCLB] = clb_pa + clb_i * HBA_CLB_SIZE;
        port_regs[HBA_REG_PxFB] = fis_pa + fis_i * HBA_FIS_SIZE;

        *port = (hba_port_t){
            .regs = port_regs,
            .ssts = port_regs[HBA_REG_PxSSTS],
            .cmdlst = (hba_cmdh_t*)(clb_pg_va + clb_i * HBA_CLB_SIZE),
            .fis = (void*)(fis_pg_va + fis_i * HBA_FIS_SIZE),
        };

        // 置于就绪状态
        port_regs[HBA_REG_PxCI] = 0;
        // clear register
        port_regs[HBA_REG_PxSERR] = -1;

        hba.ports[i] = port;

        if (!HBA_PxSSTS_SPD(port->ssts)) {
            continue;
        }

        wait_until(!(port_regs[HBA_REG_PxCMD] & HBA_PxCMD_CR));

        port_regs[HBA_REG_PxCMD] |= HBA_PxCMD_FRE;
        port_regs[HBA_REG_PxCMD] |= HBA_PxCMD_ST;

        if (!hba_init_device(port)) {
            LOGE("AHCI port %d init failed\n", i);
            vfree(hba.ports[i]);
            hba.ports[i] = NULL;
            continue;
        }

        LOGI("SATA%d: %s, sector size=%dB, sector=%u\n",
            i,
            port->device->model,
            port->device->block_size,
            (u32)port->device->max_lba
        );

        // TODO: mount disk
        ahci_register_device(port->device);
    }
}

u32 ahci_get_port_usage() {
    return hba.ports_bmap;
}

hba_port_t* ahci_get_port(u32 index) {
    if (index >= 32) {
        return 0;
    }
    return hba.ports[index];
}

int ahci_try_send(hba_port_t* port, int slot) {
    int retry = 0, mask = 1 << slot;

    wait_until(!(port->regs[HBA_REG_PxTFD] & (HBA_PxTFD_BSY | HBA_PxTFD_DRQ)));

    port->regs[HBA_REG_PxIS] = -1;

    while (retry < 3) {
        port->regs[HBA_REG_PxCI] = mask;

        wait_until_expire(!(port->regs[HBA_REG_PxCI] & mask), 1000000);

        port->regs[HBA_REG_PxCI] &= ~mask;
        if (port->regs[HBA_REG_PxTFD] & HBA_PxTFD_ERR) {
            sata_read_error(port);
            retry++;
        } else {
            break;
        }
    }

    port->regs[HBA_REG_PxIS] = -1;
    return retry < 3;
}

char sata_ifs[][20] = { "Not detected",
    "SATA I (1.5Gbps)",
    "SATA II (3.0Gbps)",
    "SATA III (6.0Gbps)"
};

void ahci_list_device() {
    LOGI("Version: %X; Ports: %d; Slot: %d\n",
        hba.version,
        hba.ports_num,
        hba.cmd_slots);

    hba_port_t* port;
    for (size_t i = 0; i < 32; i++) {
        port = hba.ports[i];
        if (port == NULL) {
            continue;
        }

        int device_state = HBA_PxSSTS_IPM(port->ssts);

        LOGI("\t Port %d: %s (%X)\n",
            i,
            &sata_ifs[device_state],
            port->device->flags);

        hba_dev_t* dev_info = port->device;
        if (!device_state || !dev_info) {
            continue;
        }
        LOGI("\t\t capacity: %d KiB\n",
            (dev_info->max_lba * dev_info->block_size) >> 10);
        LOGI("\t\t block size: %dB\n", dev_info->block_size);
        LOGI("\t\t block/sector: %d\n", dev_info->block_per_sec);
        LOGI("\t\t alignment: %dB\n", dev_info->alignment_offset);
        LOGI("\t\t capabilities: %X\n", dev_info->capabilities);
        LOGI("\t\t model: %s\n", &dev_info->model);
        LOGI("\t\t serial: %s\n", &dev_info->serial_num);

        // u8* buf = valloc_dma(dev_info->block_size);
        // int ret = port->device->ops.read_buffer(port->device, 0, buf, dev_info->block_size);
        // if (ret) {
        //     buf[dev_info->block_size - 1] = 0;
        //     LOGP("%s\n", buf);
        // }
    }
}

bool ahci_io_is_req_done(void* data) {
    hba_port_t* port = (hba_port_t*)data;
    if ((port->regs[HBA_REG_PxTFD] & (HBA_PxTFD_BSY | HBA_PxTFD_DRQ))) {
        return false;
    }
    if ((port->regs[HBA_REG_PxCI] ^ port->cmdctx.tracked_ci) && hba.base[HBA_REG_IS]) {
        return false;
    }
    return true;
}

void ahci_post(hba_port_t* port, hba_cmds_t* state, int slot) {
    int bitmask = 1 << slot;

    blkio_req_t* req = (blkio_req_t*)state->state_ctx;
    req->is_req_done = ahci_io_is_req_done;
    req->track_data = port;

    // 确保端口是空闲的
    wait_until(!(port->regs[HBA_REG_PxTFD] & (HBA_PxTFD_BSY | HBA_PxTFD_DRQ)));

    hba_clear_reg(port->regs[HBA_REG_PxIS]);

    port->cmdctx.issued[slot] = state;
    port->cmdctx.tracked_ci |= bitmask;
    port->regs[HBA_REG_PxCI] |= bitmask;
}